package main

import (
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
)

func RunServer() error {
	server, err := net.Listen(settings.Protocol, fmt.Sprintf("%s:%d", settings.Ip, settings.Port))
	if err != nil {
		return err
	}

	tip := fmt.Sprintf("Start listening(%s:%d) ... ", settings.Ip, settings.Port)
	fmt.Println(tip)

	for {
		client, err := server.Accept()
		if err != nil {
			fmt.Println(err)
			continue
		}
		go HandleClient(client)
	}
}

func HandleClient(conn net.Conn) {
	defer conn.Close()

	// 核对协议及验证身份
	if err := HandleAuth(conn); err != nil {
		fmt.Println(err)
		return
	}

	// 处理客户端请求
	if err := HandleClientRequest(conn); err != nil {
		fmt.Println(err)
		return
	}
}

func HandleAuth(conn net.Conn) error {
	packet := UnpackHandshakePacket(conn)

	if packet.Err != nil {
		return packet.Err
	}

	for i := 0; i < int(packet.NMethods); i++ {
		// 目前只处理不需要认证的情况
		if packet.Methods[i] == 0x00 {
			echo := EchoHandshakePacket{SOCKS5_VERSION, 0x00}
			_, err := conn.Write(echo.bytes())
			return err
		}
	}

	return errors.New("authentication failed")
}

func HandleClientRequest(conn net.Conn) error {
	packet := UnpackRequestPacket(conn)

	if packet.Err != nil {
		return packet.Err
	}

	fmt.Println(packet.ReqDesc())

	if packet.Cmd == SOCKS5_CMD_CONNECT {
		remote, err := net.Dial("tcp", packet.AddrAndPortString())
		if err != nil {
			reply := ReplyPacket{Ver: SOCKS5_VERSION, Rep: SOCKS5_REP_HOST_UNREACHABLE, Atyp: packet.Atyp}
			conn.Write(reply.bytes())
			return err
		}
		defer remote.Close()

		reply := ReplyPacket{Ver: SOCKS5_VERSION, Rep: SOCKS5_REP_SUCCESS, Atyp: packet.Atyp}
		conn.Write(reply.bytes())

		wg := new(sync.WaitGroup)
		wg.Add(2)

		go func() {
			defer wg.Done()
			io.Copy(conn, remote)
		}()

		go func() {
			defer wg.Done()
			io.Copy(remote, conn)
		}()

		wg.Wait()

	}

	return nil
}
